#!/usr/bin/env bash
# Author: Noel Chalmers

# set -x #echo on

# #################################################
# helper functions
# #################################################
function display_help()
{
  echo "rocHPL MPI run helper script"
  echo "./mpirun_rochpl "
  echo "    [-P]    Specific MPI grid size: the number of         "
  echo "            rows in MPI grid.                             "
  echo "    [-Q]    Specific MPI grid size: the number of         "
  echo "            columns in MPI grid.                          "
  echo "    [-p]    Specific node-local MPI grid size: the number "
  echo "            of rows in node-local MPI grid. Must evenly   "
  echo "            divide P.                                     "
  echo "    [-q]    Specific node-local MPI grid size: the number "
  echo "            of columns in node-local MPI grid. Must evenly"
  echo "            divide Q.                                     "
  echo "    [-N]    Specific matrix size: the number of           "
  echo "            rows/columns in global matrix.                "
  echo "    [--NB]  Specific panel size: the number of            "
  echo "            rows/columns in panels.                       "
  echo "    [-f]    Specific split fraction: the percentange to   "
  echo "            split the trailing submatrix.                 "
  echo "    [-i]    Input file. When set, all other commnand      "
  echo "            line parameters are ignored, and problem      "
  echo "            parameters are read from input file.          "
  echo "    [-h|--help] prints this help message                  "
  echo "    [--version] Print rocHPL version number.              "
}

# This function is helpful for dockerfiles that do not have sudo installed, but the default user is root
# true is a system command that completes successfully, function returns success
# prereq: ${ID} must be defined before calling
supported_distro( )
{
  if [ -z ${ID+foo} ]; then
    printf "supported_distro(): \$ID must be set\n"
    exit 2
  fi

  case "${ID}" in
    ubuntu|centos|rhel|fedora|sles)
        true
        ;;
    *)  printf "This script is currently supported on Ubuntu, CentOS, RHEL, Fedora and SLES\n"
        exit 2
        ;;
  esac
}

# #################################################
# Pre-requisites check
# #################################################
# Exit code 0: alls well
# Exit code 1: problems with getopt
# Exit code 2: problems with supported platforms

# check if getopt command is installed
type getopt > /dev/null
if [[ $? -ne 0 ]]; then
  echo "This script uses getopt to parse arguments; try installing the util-linux package";
  exit 1
fi

# os-release file describes the system
if [[ -e "/etc/os-release" ]]; then
  source /etc/os-release
else
  echo "This script depends on the /etc/os-release file"
  exit 2
fi

# The following function exits script if an unsupported distro is detected
supported_distro

# #################################################
# global variables
# #################################################
# Grab options from CMake config
rochpl_bin=@CMAKE_INSTALL_PREFIX@/bin/rochpl
mpi_bin=@MPIEXEC_EXECUTABLE@
rochpl_runscript=$(dirname "$0")/run_rochpl #assume run_rochpl is in the same location

P=1
Q=1
p=-1
q=-1
N=45312
NB=384
frac=0.6

filename=HPL.dat
inputfile=false
cmdrun=false

# #################################################
# MPI Args
# #################################################
# count the number of physical cores
num_cpu_cores=$(lscpu | grep "Core(s)" | awk '{print $4}')
num_cpu_sockets=$(lscpu | grep Socket | awk '{print $2}')
total_cpu_cores=$(($num_cpu_cores*$num_cpu_sockets))

#Default MPI options
mpi_args=

#Check if using OpenMPI
if [[ $(${mpi_bin} --version | grep "open-mpi") ]]; then
  mpi_args+=" --map-by node --rank-by node --bind-to none "

  #Check if this is OpenMPI+UCX
  ompi_info=$(dirname ${mpi_bin})/ompi_info
  if [[ $(${ompi_info} | grep "MCA pml: ucx") ]]; then
    # ucx-specific args
    mpi_args="--mca pml ucx --mca btl ^vader,tcp,openib,uct ${mpi_args}"
  fi
fi

# #################################################
# Parameter parsing
# #################################################

# check if we have a modern version of getopt that can handle whitespace and long parameters
getopt -T
if [[ $? -eq 4 ]]; then
  GETOPT_PARSE=$(getopt --name "${0}" --longoptions NB:,help,version, --options hP:Q:p:q:N:i:f: -- "$@")
else
  echo "Need a new version of getopt"
  exit 1
fi

if [[ $? -ne 0 ]]; then
  echo "getopt invocation failed; could not parse the command line";
  exit 1
fi

eval set -- "${GETOPT_PARSE}"

while true; do
  case "${1}" in
    -h|--help)
        display_help
        exit 0
        ;;
    --version)
        ${mpi_bin} -np 1 ${mpi_args} ${rochpl_runscript} --version
        exit 0
        ;;
    -P)
        P=${2}
        shift 2 ;;
    -Q)
        Q=${2}
        shift 2 ;;
    -p)
        p=${2}
        shift 2 ;;
    -q)
        q=${2}
        shift 2 ;;
    -N)
        N=${2}
        cmdrun=true
        shift 2 ;;
    --NB)
        NB=${2}
        cmdrun=true
        shift 2 ;;
    -f)
        frac=${2}
        shift 2 ;;
    -i)
        filename=${2}
        inputfile=true
        shift 2 ;;
    --) shift ; break ;;
    *)  echo "Unexpected command line parameter received; aborting";
        exit 1
        ;;
  esac
done

#if nothing but np and ppn parameters where given, default to running
# with default input file
if [[ "${inputfile}" == false && "${cmdrun}" == false ]]; then
  inputfile=true
fi

np=$(($P*$Q))
if [[ "$np" -lt 1 ]]; then
  echo "Invalid MPI grid parameters; aborting";
  exit 1
fi

if [[ "${inputfile}" == true ]]; then
  rochpl_args="-P ${P} -Q ${Q} -p ${p} -q ${q} -i ${filename} -f ${frac}"
else
  rochpl_args="-P ${P} -Q ${Q} -p ${p} -q ${q} -N ${N} --NB ${NB} -f ${frac}"
fi

#run
${mpi_bin} -np ${np} ${mpi_args} ${rochpl_runscript} ${rochpl_args}
